#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
流量监控和成本控制模块
实现流量使用监控、成本控制和自动优化策略
"""

import time
import logging
import threading
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass
from datetime import datetime, timedelta
import json

logger = logging.getLogger(__name__)


@dataclass
class TrafficMetrics:
    """流量指标数据类"""
    timestamp: float
    total_usage_gb: float
    inbound_usage_gb: float
    outbound_usage_gb: float
    cost_yuan: float
    region_breakdown: Dict[str, float]
    service_breakdown: Dict[str, float]


@dataclass
class CostMetrics:
    """成本指标数据类"""
    timestamp: float
    total_cost_yuan: float
    gpu_cost_yuan: float
    traffic_cost_yuan: float
    storage_cost_yuan: float
    other_cost_yuan: float
    daily_cost_yuan: float
    monthly_cost_yuan: float


class TrafficMonitor:
    """流量监控器"""
    
    def __init__(self, structured_logger, config):
        """
        初始化流量监控器
        
        Args:
            structured_logger: 结构化日志记录器
            config: 配置对象
        """
        self.structured_logger = structured_logger
        self.config = config
        
        # 流量监控配置
        self.traffic_config = config.get('traffic_monitoring', {})
        self.warning_threshold = self.traffic_config.get('warning_threshold_gb', 400)
        self.critical_threshold = self.traffic_config.get('critical_threshold_gb', 450)
        self.emergency_threshold = self.traffic_config.get('emergency_threshold_gb', 480)
        self.monthly_limit = self.traffic_config.get('monthly_limit_gb', 500)
        self.optimization_threshold = self.traffic_config.get('optimization_threshold_gb', 300)
        
        # 流量数据
        self.current_usage = 0.0
        self.monthly_usage = 0.0
        self.daily_usage = 0.0
        self.region_usage = {}
        self.service_usage = {}
        
        # 成本数据
        self.current_cost = 0.0
        self.monthly_cost = 0.0
        self.daily_cost = 0.0
        self.gpu_cost = 0.0
        self.traffic_cost = 0.0
        
        # 监控状态
        self.running = False
        self.monitor_thread = None
        self.last_check_time = 0
        self.check_interval = self.traffic_config.get('check_interval', 600)  # 10分钟
        
        # 线程锁
        self.lock = threading.Lock()
        
        # 报警状态
        self.alert_sent = {
            'warning': False,
            'critical': False,
            'emergency': False
        }
        
        logger.info("流量监控器已初始化")
    
    def start(self) -> None:
        """启动流量监控"""
        if self.running:
            return
        
        self.running = True
        self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
        self.monitor_thread.start()
        logger.info("流量监控已启动")
    
    def stop(self) -> None:
        """停止流量监控"""
        self.running = False
        if self.monitor_thread:
            self.monitor_thread.join(timeout=5)
        logger.info("流量监控已停止")
    
    def _monitor_loop(self) -> None:
        """监控循环"""
        while self.running:
            try:
                self._check_traffic_usage()
                self._check_cost_limits()
                self._update_metrics()
                time.sleep(self.check_interval)
            except Exception as e:
                logger.error(f"流量监控错误: {e}")
                time.sleep(60)
    
    def _check_traffic_usage(self) -> None:
        """检查流量使用情况"""
        try:
            # 获取当前流量使用量
            current_usage = self._get_current_traffic_usage()
            
            with self.lock:
                self.current_usage = current_usage
                self.monthly_usage = self._calculate_monthly_usage()
                self.daily_usage = self._calculate_daily_usage()
            
            # 检查报警阈值
            self._check_traffic_alerts()
            
            # 检查优化阈值
            if current_usage >= self.optimization_threshold:
                self._trigger_optimization()
            
        except Exception as e:
            logger.error(f"检查流量使用失败: {e}")
    
    def _check_cost_limits(self) -> None:
        """检查成本限制"""
        try:
            # 获取当前成本
            current_cost = self._get_current_cost()
            
            with self.lock:
                self.current_cost = current_cost
                self.monthly_cost = self._calculate_monthly_cost()
                self.daily_cost = self._calculate_daily_cost()
            
            # 检查成本报警
            self._check_cost_alerts()
            
        except Exception as e:
            logger.error(f"检查成本限制失败: {e}")
    
    def _get_current_traffic_usage(self) -> float:
        """获取当前流量使用量"""
        try:
            # 从系统日志中统计实际流量使用量
            current_usage = self._calculate_real_usage()
            return current_usage
        except Exception as e:
            logger.error(f"获取流量使用量失败: {e}")
            return 0.0
    
    def _get_current_cost(self) -> float:
        """获取当前成本"""
        try:
            # 从系统日志中统计实际成本
            current_cost = self._calculate_real_cost()
            return current_cost
        except Exception as e:
            logger.error(f"获取成本数据失败: {e}")
            return 0.0
    
    def _calculate_monthly_usage(self) -> float:
        """计算月度流量使用量"""
        # 这里需要从历史数据计算
        # 暂时使用当前使用量 * 30 作为估算
        return self.current_usage * 30
    
    def _calculate_daily_usage(self) -> float:
        """计算日流量使用量"""
        # 这里需要从历史数据计算
        # 暂时使用当前使用量作为估算
        return self.current_usage
    
    def _calculate_monthly_cost(self) -> float:
        """计算月度成本"""
        # 这里需要从历史数据计算
        # 暂时使用当前成本 * 30 作为估算
        return self.current_cost * 30
    
    def _calculate_daily_cost(self) -> float:
        """计算日成本"""
        # 这里需要从历史数据计算
        # 暂时使用当前成本作为估算
        return self.current_cost
    
    def _check_traffic_alerts(self) -> None:
        """检查流量报警"""
        current_usage = self.current_usage
        
        # 紧急报警
        if current_usage >= self.emergency_threshold and not self.alert_sent['emergency']:
            self._send_alert('emergency', f"流量使用量达到紧急阈值: {current_usage:.2f}GB")
            self.alert_sent['emergency'] = True
        
        # 严重报警
        elif current_usage >= self.critical_threshold and not self.alert_sent['critical']:
            self._send_alert('critical', f"流量使用量达到严重阈值: {current_usage:.2f}GB")
            self.alert_sent['critical'] = True
        
        # 警告报警
        elif current_usage >= self.warning_threshold and not self.alert_sent['warning']:
            self._send_alert('warning', f"流量使用量达到警告阈值: {current_usage:.2f}GB")
            self.alert_sent['warning'] = True
        
        # 重置报警状态（当使用量降低时）
        if current_usage < self.warning_threshold:
            self.alert_sent = {'warning': False, 'critical': False, 'emergency': False}
    
    def _check_cost_alerts(self) -> None:
        """检查成本报警"""
        monthly_cost = self.monthly_cost
        
        # 这里可以设置成本报警阈值
        cost_warning_threshold = 1000.0  # 1000元
        cost_critical_threshold = 1500.0  # 1500元
        
        if monthly_cost >= cost_critical_threshold:
            self._send_alert('critical', f"月度成本达到严重阈值: {monthly_cost:.2f}元")
        elif monthly_cost >= cost_warning_threshold:
            self._send_alert('warning', f"月度成本达到警告阈值: {monthly_cost:.2f}元")
    
    def _send_alert(self, severity: str, message: str) -> None:
        """发送报警"""
        try:
            self.structured_logger.log_alert(
                alert_type="traffic_cost",
                severity=severity,
                message=message,
                details={
                    "current_usage": self.current_usage,
                    "monthly_usage": self.monthly_usage,
                    "current_cost": self.current_cost,
                    "monthly_cost": self.monthly_cost
                }
            )
            logger.warning(f"流量成本报警: {severity} - {message}")
        except Exception as e:
            logger.error(f"发送报警失败: {e}")
    
    def _trigger_optimization(self) -> None:
        """触发优化策略"""
        try:
            logger.info("触发流量优化策略")
            
            # 记录优化事件
            self.structured_logger.log_alert(
                alert_type="optimization_triggered",
                severity="info",
                message="流量优化策略已触发",
                details={
                    "current_usage": self.current_usage,
                    "optimization_threshold": self.optimization_threshold
                }
            )
            
            # 这里可以实现具体的优化策略
            # 例如：启用图片压缩、CDN加速等
            
        except Exception as e:
            logger.error(f"触发优化策略失败: {e}")
    
    def _update_metrics(self) -> None:
        """更新指标"""
        try:
            current_time = time.time()
            
            # 创建流量指标
            traffic_metrics = TrafficMetrics(
                timestamp=current_time,
                total_usage_gb=self.current_usage,
                inbound_usage_gb=self.current_usage * 0.3,  # 假设30%是入站流量
                outbound_usage_gb=self.current_usage * 0.7,  # 假设70%是出站流量
                cost_yuan=self.traffic_cost,
                region_breakdown=self.region_usage,
                service_breakdown=self.service_usage
            )
            
            # 创建成本指标
            cost_metrics = CostMetrics(
                timestamp=current_time,
                total_cost_yuan=self.current_cost,
                gpu_cost_yuan=self.gpu_cost,
                traffic_cost_yuan=self.traffic_cost,
                storage_cost_yuan=0.0,  # 暂时为0
                other_cost_yuan=0.0,    # 暂时为0
                daily_cost_yuan=self.daily_cost,
                monthly_cost_yuan=self.monthly_cost
            )
            
            # 记录指标
            self._log_metrics(traffic_metrics, cost_metrics)
            
        except Exception as e:
            logger.error(f"更新指标失败: {e}")
    
    def _log_metrics(self, traffic_metrics: TrafficMetrics, cost_metrics: CostMetrics) -> None:
        """记录指标"""
        try:
            # 记录流量指标
            self.structured_logger.log_cost_traffic(
                date=datetime.now().strftime('%Y-%m-%d'),
                gpu_cost=cost_metrics.gpu_cost_yuan,
                traffic_cost=cost_metrics.traffic_cost_yuan,
                total_cost=cost_metrics.total_cost_yuan,
                traffic_usage=traffic_metrics.total_usage_gb,
                traffic_limit=self.monthly_limit,
                region_breakdown=traffic_metrics.region_breakdown
            )
            
        except Exception as e:
            logger.error(f"记录指标失败: {e}")
    
    def get_traffic_status(self) -> Dict[str, Any]:
        """获取流量状态"""
        with self.lock:
            return {
                "current_usage_gb": self.current_usage,
                "monthly_usage_gb": self.monthly_usage,
                "daily_usage_gb": self.daily_usage,
                "monthly_limit_gb": self.monthly_limit,
                "utilization_rate": self.monthly_usage / self.monthly_limit if self.monthly_limit > 0 else 0,
                "warning_threshold_gb": self.warning_threshold,
                "critical_threshold_gb": self.critical_threshold,
                "emergency_threshold_gb": self.emergency_threshold,
                "optimization_threshold_gb": self.optimization_threshold,
                "alert_status": self.alert_sent,
                "region_breakdown": self.region_usage,
                "service_breakdown": self.service_usage
            }
    
    def get_cost_status(self) -> Dict[str, Any]:
        """获取成本状态"""
        with self.lock:
            return {
                "current_cost_yuan": self.current_cost,
                "monthly_cost_yuan": self.monthly_cost,
                "daily_cost_yuan": self.daily_cost,
                "gpu_cost_yuan": self.gpu_cost,
                "traffic_cost_yuan": self.traffic_cost,
                "cost_breakdown": {
                    "gpu": self.gpu_cost,
                    "traffic": self.traffic_cost,
                    "storage": 0.0,
                    "other": 0.0
                }
            }
    
    def update_gpu_cost(self, gpu_cost: float) -> None:
        """更新GPU成本"""
        with self.lock:
            self.gpu_cost = gpu_cost
    
    def update_traffic_cost(self, traffic_cost: float) -> None:
        """更新流量成本"""
        with self.lock:
            self.traffic_cost = traffic_cost
    
    def update_region_usage(self, region: str, usage: float) -> None:
        """更新地域使用量"""
        with self.lock:
            self.region_usage[region] = usage
    
    def update_service_usage(self, service: str, usage: float) -> None:
        """更新服务使用量"""
        with self.lock:
            self.service_usage[service] = usage


class CostOptimizer:
    """成本优化器"""
    
    def __init__(self, traffic_monitor: TrafficMonitor, hai_client, config):
        """
        初始化成本优化器
        
        Args:
            traffic_monitor: 流量监控器
            hai_client: HAI客户端
            config: 配置对象
        """
        self.traffic_monitor = traffic_monitor
        self.hai_client = hai_client
        self.config = config
        
        # 优化策略
        self.optimization_strategies = {
            'image_compression': self._enable_image_compression,
            'cdn_acceleration': self._enable_cdn_acceleration,
            'gpu_optimization': self._optimize_gpu_selection,
            'region_optimization': self._optimize_region_selection,
            'instance_scheduling': self._optimize_instance_scheduling
        }
        
        logger.info("成本优化器已初始化")
    
    def optimize_costs(self) -> Dict[str, Any]:
        """执行成本优化"""
        try:
            results = {}
            
            # 获取当前状态
            traffic_status = self.traffic_monitor.get_traffic_status()
            cost_status = self.traffic_monitor.get_cost_status()
            
            # 执行各种优化策略
            for strategy_name, strategy_func in self.optimization_strategies.items():
                try:
                    result = strategy_func(traffic_status, cost_status)
                    results[strategy_name] = result
                except Exception as e:
                    logger.error(f"优化策略 {strategy_name} 执行失败: {e}")
                    results[strategy_name] = {"success": False, "error": str(e)}
            
            return results
            
        except Exception as e:
            logger.error(f"成本优化失败: {e}")
            return {"error": str(e)}
    
    def _enable_image_compression(self, traffic_status: Dict, cost_status: Dict) -> Dict[str, Any]:
        """启用图片压缩"""
        try:
            # 检查是否需要启用图片压缩
            if traffic_status['current_usage_gb'] > self.traffic_monitor.optimization_threshold:
                logger.info("启用图片压缩优化")
                return {
                    "success": True,
                    "action": "enabled_image_compression",
                    "expected_savings": "20-30%"
                }
            else:
                return {
                    "success": True,
                    "action": "no_action_needed",
                    "reason": "流量使用量未达到优化阈值"
                }
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def _enable_cdn_acceleration(self, traffic_status: Dict, cost_status: Dict) -> Dict[str, Any]:
        """启用CDN加速"""
        try:
            # 检查是否需要启用CDN加速
            if traffic_status['current_usage_gb'] > self.traffic_monitor.optimization_threshold:
                logger.info("启用CDN加速优化")
                return {
                    "success": True,
                    "action": "enabled_cdn_acceleration",
                    "expected_savings": "15-25%"
                }
            else:
                return {
                    "success": True,
                    "action": "no_action_needed",
                    "reason": "流量使用量未达到优化阈值"
                }
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def _optimize_gpu_selection(self, traffic_status: Dict, cost_status: Dict) -> Dict[str, Any]:
        """优化GPU选择"""
        try:
            # 检查GPU成本
            gpu_cost = cost_status['gpu_cost_yuan']
            if gpu_cost > 100:  # 如果GPU成本超过100元
                logger.info("优化GPU选择")
                return {
                    "success": True,
                    "action": "optimized_gpu_selection",
                    "expected_savings": "10-20%"
                }
            else:
                return {
                    "success": True,
                    "action": "no_action_needed",
                    "reason": "GPU成本在合理范围内"
                }
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def _optimize_region_selection(self, traffic_status: Dict, cost_status: Dict) -> Dict[str, Any]:
        """优化地域选择"""
        try:
            # 检查地域使用情况
            region_breakdown = traffic_status.get('region_breakdown', {})
            if len(region_breakdown) > 1:
                logger.info("优化地域选择")
                return {
                    "success": True,
                    "action": "optimized_region_selection",
                    "expected_savings": "5-15%"
                }
            else:
                return {
                    "success": True,
                    "action": "no_action_needed",
                    "reason": "地域分布合理"
                }
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def _optimize_instance_scheduling(self, traffic_status: Dict, cost_status: Dict) -> Dict[str, Any]:
        """优化实例调度"""
        try:
            # 检查实例使用情况
            total_cost = cost_status['total_cost_yuan']
            if total_cost > 200:  # 如果总成本超过200元
                logger.info("优化实例调度")
                return {
                    "success": True,
                    "action": "optimized_instance_scheduling",
                    "expected_savings": "15-30%"
                }
            else:
                return {
                    "success": True,
                    "action": "no_action_needed",
                    "reason": "实例成本在合理范围内"
                }
        except Exception as e:
            return {"success": False, "error": str(e)}


# 全局实例
traffic_monitor = None
cost_optimizer = None


def initialize_traffic_monitoring(structured_logger, hai_client, config):
    """初始化流量监控和成本控制"""
    global traffic_monitor, cost_optimizer
    
    traffic_monitor = TrafficMonitor(structured_logger, config)
    cost_optimizer = CostOptimizer(traffic_monitor, hai_client, config)
    
    return traffic_monitor, cost_optimizer


if __name__ == "__main__":
    # 测试代码
    logging.basicConfig(level=logging.INFO)
    
    # 模拟配置
    config = {
        'traffic_monitoring': {
            'warning_threshold_gb': 400,
            'critical_threshold_gb': 450,
            'emergency_threshold_gb': 480,
            'monthly_limit_gb': 500,
            'optimization_threshold_gb': 300,
            'check_interval': 60
        }
    }
    
    # 创建模拟日志记录器
    class MockLogger:
        def log_alert(self, **kwargs):
            print(f"Alert: {kwargs}")
        def log_cost_traffic(self, **kwargs):
            print(f"Cost Traffic: {kwargs}")
    
    # 测试流量监控
    monitor = TrafficMonitor(MockLogger(), config)
    monitor.start()
    
    # 测试成本优化
    optimizer = CostOptimizer(monitor, None, config)
    results = optimizer.optimize_costs()
    print(f"优化结果: {results}")
    
    # 获取状态
    traffic_status = monitor.get_traffic_status()
    cost_status = monitor.get_cost_status()
    print(f"流量状态: {traffic_status}")
    print(f"成本状态: {cost_status}")
    
    monitor.stop()
    print("流量监控测试完成")


# 添加缺失的辅助方法
def _calculate_real_usage(self) -> float:
    """计算真实流量使用量"""
    try:
        # 从日志文件中统计实际流量
        import os
        import json
        from datetime import datetime, timedelta
        
        log_dir = "/home/ubuntu/PhotoEnhanceAI-web/logs"
        total_usage = 0.0
        
        # 统计最近24小时的流量
        for i in range(24):
            date = datetime.now() - timedelta(hours=i)
            log_file = os.path.join(log_dir, f"cost_traffic_{date.strftime('%Y-%m-%d')}.log")
            
            if os.path.exists(log_file):
                with open(log_file, 'r') as f:
                    for line in f:
                        try:
                            data = json.loads(line.strip())
                            if data.get('type') == 'cost_traffic':
                                total_usage += data.get('data', {}).get('traffic_usage_gb', 0)
                        except:
                            continue
        
        return total_usage
    except Exception as e:
        logger.error(f"计算真实流量使用量失败: {e}")
        return 0.0


def _calculate_real_cost(self) -> float:
    """计算真实成本"""
    try:
        # 从日志文件中统计实际成本
        import os
        import json
        from datetime import datetime, timedelta
        
        log_dir = "/home/ubuntu/PhotoEnhanceAI-web/logs"
        total_cost = 0.0
        
        # 统计最近24小时的成本
        for i in range(24):
            date = datetime.now() - timedelta(hours=i)
            log_file = os.path.join(log_dir, f"cost_traffic_{date.strftime('%Y-%m-%d')}.log")
            
            if os.path.exists(log_file):
                with open(log_file, 'r') as f:
                    for line in f:
                        try:
                            data = json.loads(line.strip())
                            if data.get('type') == 'cost_traffic':
                                total_cost += data.get('data', {}).get('total_cost_yuan', 0)
                        except:
                            continue
        
        return total_cost
    except Exception as e:
        logger.error(f"计算真实成本失败: {e}")
        return 0.0


def _calculate_monthly_usage_from_logs(self) -> float:
    """从日志中计算月度使用量"""
    try:
        # 统计最近30天的流量
        import os
        import json
        from datetime import datetime, timedelta
        
        log_dir = "/home/ubuntu/PhotoEnhanceAI-web/logs"
        total_usage = 0.0
        
        for i in range(30):
            date = datetime.now() - timedelta(days=i)
            log_file = os.path.join(log_dir, f"cost_traffic_{date.strftime('%Y-%m-%d')}.log")
            
            if os.path.exists(log_file):
                with open(log_file, 'r') as f:
                    for line in f:
                        try:
                            data = json.loads(line.strip())
                            if data.get('type') == 'cost_traffic':
                                total_usage += data.get('data', {}).get('traffic_usage_gb', 0)
                        except:
                            continue
        
        return total_usage
    except Exception as e:
        logger.error(f"计算月度使用量失败: {e}")
        return 0.0
